#from openopt import NLP, GLP
from scipy import optimize
#from numpy import cos, arange, ones, asarray, abs, zeros, sqrt, asscalar, array, pi
import numpy as np
import pylab
import JT
import csv
import JT.JTconfig as config
import os
from JT import UploadToGoogle
import my_auth
from socket import gethostname
from datetime import datetime
import time
import JT.PlotResultsFromFile
user = gethostname()

from win32api import GetSystemMetrics

class Test():
    def __init__(self,logging=True,upload=True):            

#===============================================================================
#    Google Data Upload Setup
#===============================================================================
        cred = my_auth.credentials()
        username = cred.name
        password = cred.pw
        self._connectSuccess = False
        
        if upload:
            try:
                self._uploader = UploadToGoogle.Uploader(username,password)
                # Choose newest sheet and book
                self._uploader._run(book='0',sheet='0')
                self._connectSuccess = True
            except:
                print 'Could not connect to Google'
        
        if logging:
            logTag = 'GD_Results'     #This tag can be used to identify trials
            logName = time.strftime('%Y%m%d_%H%M%S')+'_'+logTag+'.log'
            logDir = 'C:/Users/bryan/bryan-code/trunk/Results'
            print 'Final Result Logging dir: ', logDir 
            self._logFile = csv.writer(open(os.path.join(logDir,logName),'wt'),
                                   dialect='excel-tab')
            # Write Headers
            self._logFile.writerow(('Iteration','Cost','Starting Pose0(mm)',
                                    'Starting Pose1(mm)','Starting Pose2(mm)',
                                    'Starting Pose3(deg)','Starting Pose4(deg)',
                                    'Starting Pose5(deg)','Final Pose0(mm)',
                                    'Final Pose1(mm)','Final Pose2(mm)',
                                    'Final Pose3(deg)','Final Pose4(deg)',
                                    'Final Pose5(deg)'))
            
#===============================================================================
#    Setup DRR
#===============================================================================
        xraycam = JT.XRayCamera()
        drr = JT.DRR()
        drr.SetXRayCamera(xraycam)
        drr.SetBlackOnWhiteImage()
        drr.InteractionOff()
        
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(config.calFileLat)
        cal.LoadStaFile(config.staFile)
        
        xraycam.SetExternalCalibration(cal)
        
        self.volume = JT.Volume()
        self.volume.SetVolumeFilename(config.inputVolFileName)
        self.volume.SetOriginalTransform(cal._VolumeT)
        self.volume.SetBodyCenter((3.906, 2.17, -7.378))
        #self.volume.UseRayCastMapper()
        drr.AddVolume(self.volume)
        
        drr._renWin.SetOffScreenRendering(1)
        
        # Set Fixed image using path
        self.fixedImage = JT.FixedImage()
        self.fixedImage.SetFileName(config.fixedImageName)

        self.reg=JT.Registration()
        self.reg.SetFixedImage(self.fixedImage.GetImage(0))
        self.reg.SetMovingImageGenerator(drr)
        self.reg.SetRegionOfInterest(config.roi)
        
        self.Testing = JT.Testing()


    def run_test(self,n_pose=50,level=4,method='gc',debug=True,plot=False,upload_data=True):    
    #===============================================================================
    #    Setup Optimizer
    #===============================================================================
    # Select the desired starting pose, between [0,200)
        
        # set Metric
        self.reg.SetImageMetric(method)
        
        from JT.ITKOptimizers import RegularStepGradientDescentOptimizer
        gd = RegularStepGradientDescentOptimizer()
        
        # Set the starting pose
        pose = self.Testing.GetStartPose(n_pose)
        gd.SetInitialPosition(pose)
        print "Starting Pose: ", pose
        
        # Set Scales 1/1000 for translations, 1 for rotations
        scales = [1.0,1.0,1.0,500.0,500.0,500.0]
        gd.SetScales(scales)
        
        # Set the cost function
        gd.SetCostFunction(self.reg)
        
        # Modify Optimizer Parameters
        gd.SetMaximumStepLength(2)
        gd.SetMinimumStepLength(1e-1)
        gd.SetRelaxationFactor(0.75)
        gd.SetNumberOfIterations(50)
        
        # Start the optimizer
        gd._debug = True
        gd.StartOptimization()
        
        # Print Final State    
        gd.PrintSelf()
        
        # Save registration data to file for analysis of tuning parameters
        # Log status/results to a file that can be plotted later.
        initial_pose = gd._InitialPosition
        final_pose = gd._CurrentPosition
        final_cost = gd._Value
        dt = datetime.fromtimestamp(time.time())
        time_stamp = dt.strftime('%m/%d/%Y_%H:%M:%S')
        user = gethostname()
        
        my_tupel = [n_pose, final_cost, initial_pose[0], initial_pose[1],
                    initial_pose[2],np.degrees(initial_pose[3]), 
                    np.degrees(initial_pose[4]),np.degrees(initial_pose[5]),
                    final_pose[0],final_pose[1],final_pose[2],
                    np.degrees(final_pose[3]),np.degrees(final_pose[4]),
                    np.degrees(final_pose[5])]
        
        self._logFile.writerow(my_tupel)
        
        
        if upload_data and self._connectSuccess:
            output_dict = {}
            output_dict['user'] = user
            output_dict['time'] = time_stamp
            output_dict['iteration'] = str(n_pose)
            output_dict['cost'] = str(final_cost)
            output_dict['startingpose0'] = str(initial_pose[0])
            output_dict['startingpose1'] = str(initial_pose[1])
            output_dict['startingpose2'] = str(initial_pose[2])
            output_dict['startingpose3'] = str(initial_pose[3])
            output_dict['startingpose4'] = str(initial_pose[4])
            output_dict['startingpose5'] = str(initial_pose[5])
            output_dict['finalpose0'] = str(final_pose[0])
            output_dict['finalpose1'] = str(final_pose[1])
            output_dict['finalpose2'] = str(final_pose[2])
            output_dict['finalpose3'] = str(final_pose[3])
            output_dict['finalpose4'] = str(final_pose[4])
            output_dict['finalpose5'] = str(final_pose[5])
            
            self._uploader.Write(output_dict)
            #print "Data written to google spreadsheet!"
            #print record.content

if __name__ == "__main__":
    T = Test()
    for i in range(200):
        T.run_test(method='bpc',n_pose=i,plot=False)
        # Reset the reg class for the next run
        T.reg.Reset()
        #JT.PlotResultsFromFile.run()
    
    #pylab.show()


